//
// Created by twetec on 25-8-16.
//

#include "../include/network_prober/probe_detect/handle_postgresql_probe_detect.h"

#include <boost/json.hpp>
#include <postgresql/libpq-fe.h>
#include <postgresql/pg_config_ext.h>
#include <postgresql/postgres_ext.h>
#include <future>

using namespace std::literals::chrono_literals;


namespace network_prober::probe_detect
{
void tag_invoke(boost::json::value_from_tag, boost::json::value &jv, PostgreSQLProbeDetectResultIPv4 const &c)
{
    auto& ja = jv.emplace_object();

    ja["interface"] = c.interface;
    ja["target"] = c.target;
    ja["port"] = c.port;
    ja["delay"] = c.delay;
}

HandlePostgreSQLProbeDetectIPv4::HandlePostgreSQLProbeDetectIPv4(
    std::string_view interface,
    std::string_view target,
    std::uint32_t port,
    std::string_view userName,
    std::string_view password,
    std::string_view databaseName,
    std::string_view queryStatement
)
    : m_interface(interface), m_target(target), m_port(port), m_userName(userName),
    m_password(password), m_databaseName(databaseName), m_queryStatement(queryStatement)
{
}

void HandlePostgreSQLProbeDetectIPv4::detect()
{
    // 建立连接时，开始计时
    m_detectStartTime = std::chrono::steady_clock::now();

    // 指向当前对象的shared_ptr
    auto self{shared_from_this()};

    // 发起探测
    m_onDetectedCallbackFuture = std::async(std::launch::async, [this, self] {
        // 探测结果
        PostgreSQLProbeDetectResultIPv4 result;
        result.interface = m_interface;
        result.target = m_target;
        result.port = m_port;

        // 初始化数据库连接配置
        const char* paramKeys[] = {
            "host", // 要连接的主机的数字 IP 地址。 这应该是标准 IPv4 地址格式，例如 172.28.40.9。 如果您的机器支持 IPv6，您也可以使用这些地址
            "port", // 服务器主机上连接的端口号
            "dbname",   // 数据库名称
            "user", // 要连接的 PostgreSQL 用户名
            "password", // 服务器要求密码验证时使用的密码
            "connect_timeout",   // 连接时等待的最长时间，以秒为单位（写为十进制整数，例如 10）。 零、负数或未指定意味着无限期等待。 允许的最小超时为 2 秒，因此值 1 被解释为 2
            NULL
        };
        const char* paramValues[] = {
            m_target.c_str(),
            std::to_string(m_port).c_str(),
            m_databaseName.c_str(),
            m_userName.c_str(),
            m_password.c_str(),
            "5",
            NULL
        };

        // 连接PQ数据库服务器
        PGconn* connPtr = PQconnectdbParams(paramKeys, paramValues, 0);

        // 检查连接是否成功
        ConnStatusType status = PQstatus(connPtr);
        // 如果连接状态不是成功
        if (status != CONNECTION_OK) {
           // 错误信息
           SPDLOG_INFO("PgSQL观测目标连接失败：{}", PQerrorMessage(connPtr));
           // 关闭与服务器的而连接，同时释放了PGconn对象的内存 注：即使服务器连接失败，也要调用finish函数，回收内存
           PQfinish(connPtr);
           // 连接失败，延迟为-1
           result.delay = -1;
           // 回调结果
           m_onDetectedCallback(std::move(result));
           // 结束执行
           return;
        }

        // 构建SQL语句
        const char* sql = m_queryStatement.empty()? "SELECT 1" : m_queryStatement.c_str();

        // 让服务器执行SQL语句
        PGresult* execResult = PQexec(connPtr, sql);

        // 检查查询是否成功
        ExecStatusType execStatus = PQresultStatus(execResult);
        // 如果查询语句执行失败
        if (execStatus != PGRES_TUPLES_OK) {
            // 错误信息
            SPDLOG_INFO("PgSQL观测目标查询失败：{}", PQerrorMessage(connPtr));
            // 释放PGresult查询结果，有PGresult对象，就必须释放
            PQclear(execResult);
            // 关闭与服务器的而连接，同时释放了PGconn对象的内存 注：即使服务器连接失败，也要调用finish函数，回收内存
            PQfinish(connPtr);
            // 连接失败，延迟为-1
            result.delay = -1;
            // 回调结果
            m_onDetectedCallback(std::move(result));
            // 结束执行
            return;
        }

        // 释放PGresult查询结果，有PGresult对象，就必须释放
        PQclear(execResult);
        // 关闭与服务器的而连接，同时释放了PGconn对象的内存 注：即使服务器连接失败，也要调用finish函数，回收内存
        PQfinish(connPtr);

        // SQL正常执行完毕,计算延迟
        result.delay = std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::steady_clock::now() - m_detectStartTime).count();
        // 回调结果
        m_onDetectedCallback(std::move(result));
    });
}

void HandlePostgreSQLProbeDetectIPv4::setOnDetectedCallback(std::function<void(PostgreSQLProbeDetectResultIPv4&&)> callback)
{
    m_onDetectedCallback = std::move(callback);
}

HandlePostgreSQLProbeDetectIPv4::~HandlePostgreSQLProbeDetectIPv4()
{

}

void tag_invoke(boost::json::value_from_tag, boost::json::value &jv, PostgreSQLProbeDetectResultIPv6 const &c)
{
    auto& ja = jv.emplace_object();

    ja["interface"] = c.interface;
    ja["target"] = c.target;
    ja["port"] = c.port;
    ja["delay"] = c.delay;
}

HandlePostgreSQLProbeDetectIPv6::HandlePostgreSQLProbeDetectIPv6(
    std::string_view interface,
    std::string_view target,
    std::uint32_t port,
    std::string_view userName,
    std::string_view password,
    std::string_view databaseName,
    std::string_view queryStatement
)
    : m_interface(interface), m_target(target), m_port(port), m_userName(userName),
    m_password(password), m_databaseName(databaseName), m_queryStatement(queryStatement)
{
}

void HandlePostgreSQLProbeDetectIPv6::detect()
{
    // 建立连接时，开始计时
    m_detectStartTime = std::chrono::steady_clock::now();

    // 指向当前对象的shared_ptr
    auto self{shared_from_this()};

    // 发起探测
    m_onDetectedCallbackFuture = std::async(std::launch::async, [this, self] {
        // 探测结果
        PostgreSQLProbeDetectResultIPv6 result;
        result.interface = m_interface;
        result.target = m_target;
        result.port = m_port;

        // 初始化数据库连接配置
        const char* paramKeys[] = {
            "host", // 要连接的主机的数字 IP 地址。 这应该是标准 IPv4 地址格式，例如 172.28.40.9。 如果您的机器支持 IPv6，您也可以使用这些地址
            "port", // 服务器主机上连接的端口号
            "dbname",   // 数据库名称
            "user", // 要连接的 PostgreSQL 用户名
            "password", // 服务器要求密码验证时使用的密码
            "connect_timeout",   // 连接时等待的最长时间，以秒为单位（写为十进制整数，例如 10）。 零、负数或未指定意味着无限期等待。 允许的最小超时为 2 秒，因此值 1 被解释为 2
            NULL
        };
        const char* paramValues[] = {
            m_target.c_str(),
            std::to_string(m_port).c_str(),
            m_databaseName.c_str(),
            m_userName.c_str(),
            m_password.c_str(),
            "5",
            NULL
        };

        // 连接PQ数据库服务器
        PGconn* connPtr = PQconnectdbParams(paramKeys, paramValues, 0);

        // 检查连接是否成功
        ConnStatusType status = PQstatus(connPtr);
        // 如果连接状态不是成功
        if (status != CONNECTION_OK) {
            // 错误信息
            SPDLOG_INFO("PgSQL观测目标连接失败：{}", PQerrorMessage(connPtr));
            // 关闭与服务器的而连接，同时释放了PGconn对象的内存 注：即使服务器连接失败，也要调用finish函数，回收内存
            PQfinish(connPtr);
            // 连接失败，延迟为-1
            result.delay = -1;
            // 回调结果
            m_onDetectedCallback(std::move(result));
            // 结束执行
            return;
        }

        // 构建SQL语句
        const char* sql = m_queryStatement.empty()? "SELECT 1" : m_queryStatement.c_str();

        // 让服务器执行SQL语句
        PGresult* execResult = PQexec(connPtr, sql);

        // 检查查询是否成功
        ExecStatusType execStatus = PQresultStatus(execResult);
        // 如果查询语句执行失败
        if (execStatus != PGRES_TUPLES_OK) {
            // 错误信息
            SPDLOG_INFO("PgSQL观测目标查询失败：{}", PQerrorMessage(connPtr));
            // 释放PGresult查询结果，有PGresult对象，就必须释放
            PQclear(execResult);
            // 关闭与服务器的而连接，同时释放了PGconn对象的内存 注：即使服务器连接失败，也要调用finish函数，回收内存
            PQfinish(connPtr);
            // 连接失败，延迟为-1
            result.delay = -1;
            // 回调结果
            m_onDetectedCallback(std::move(result));
            // 结束执行
            return;
        }

        // 释放PGresult查询结果，有PGresult对象，就必须释放
        PQclear(execResult);
        // 关闭与服务器的而连接，同时释放了PGconn对象的内存 注：即使服务器连接失败，也要调用finish函数，回收内存
        PQfinish(connPtr);

        // SQL正常执行完毕,计算延迟
        result.delay = std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::steady_clock::now() - m_detectStartTime).count();
        // 回调结果
        m_onDetectedCallback(std::move(result));
    });
}

void HandlePostgreSQLProbeDetectIPv6::setOnDetectedCallback(std::function<void(PostgreSQLProbeDetectResultIPv6&&)> callback)
{
    m_onDetectedCallback = std::move(callback);
}

HandlePostgreSQLProbeDetectIPv6::~HandlePostgreSQLProbeDetectIPv6()
{

}
}